#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WebSocket代理桥接脚本
使用mitmproxy拦截WebSocket请求，保持原客户端与服务端的通信，
同时提供HTTP接口，允许外部系统通过HTTP请求向原客户端发送WebSocket消息。

作者: Kiro AI Assistant
版本: 1.0.0
"""

import asyncio
import json
import logging
import threading
import uuid
from dataclasses import dataclass, asdict
from datetime import datetime
from typing import Dict, List, Optional, Any
from mitmproxy import websocket,master
from mitmproxy.tools.dump import DumpMaster
from mitmproxy.options import Options
from aiohttp import web
import argparse


# ==================== 配置管理 ====================


@dataclass
class Config:
    """系统配置类，包含所有可配置的参数"""

    proxy_port: int = 8999  # mitmproxy监听端口
    api_port: int = 8998  # HTTP API服务器端口
    log_level: str = "INFO"  # 日志级别
    max_connections: int = 100  # 最大WebSocket连接数
    host: str = "127.0.0.1"  # 服务器绑定地址


# ==================== 数据模型 ====================


@dataclass
class WebSocketConnection:
    """WebSocket连接数据模型"""

    connection_id: str  # 连接唯一标识符
    client_address: str  # 客户端IP地址
    server_address: str  # 服务端地址
    created_at: datetime  # 连接建立时间
    is_active: bool = True  # 连接是否活跃
    flow: Optional[Any] = None  # mitmproxy流对象（运行时设置）

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式，用于JSON序列化"""
        return {
            "connection_id": self.connection_id,
            "client_address": self.client_address,
            "server_address": self.server_address,
            "created_at": self.created_at.isoformat(),
            "is_active": self.is_active,
        }


@dataclass
class SendMessageRequest:
    """HTTP发送消息请求数据模型"""

    connection_id: str  # 目标连接ID
    message: str  # 要发送的消息内容
    message_type: str = "text"  # 消息类型：text 或 binary

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SendMessageRequest":
        """从字典创建请求对象"""
        return cls(
            connection_id=data.get("connection_id", ""),
            message=data.get("message", ""),
            message_type=data.get("message_type", "text"),
        )

    def validate(self) -> List[str]:
        """验证请求数据，返回错误列表"""
        errors = []
        if not self.connection_id:
            errors.append("connection_id is required")
        if not self.message:
            errors.append("message is required")
        if self.message_type not in ["text", "binary"]:
            errors.append("message_type must be 'text' or 'binary'")
        return errors


# ==================== 全局配置实例 ====================

# 创建全局配置实例
config = Config()

# ==================== 日志记录和错误处理 ====================


def setup_logging():
    """配置详细的日志记录系统"""
    # 创建自定义格式化器
    formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    # 配置根日志记录器
    root_logger = logging.getLogger()
    root_logger.setLevel(getattr(logging, config.log_level.upper()))

    # 清除现有的处理器
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)

    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    root_logger.addHandler(console_handler)

    # 文件处理器 - 详细日志
    file_handler = logging.FileHandler("websocket_proxy.log", encoding="utf-8")
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(formatter)
    root_logger.addHandler(file_handler)

    # 错误日志文件处理器
    error_handler = logging.FileHandler("websocket_proxy_errors.log", encoding="utf-8")
    error_handler.setLevel(logging.ERROR)
    error_handler.setFormatter(formatter)
    root_logger.addHandler(error_handler)

    # 禁用第三方库的详细日志
    logging.getLogger("aiohttp").setLevel(logging.WARNING)
    logging.getLogger("mitmproxy").setLevel(logging.WARNING)


class ErrorHandler:
    """统一错误处理类"""

    @staticmethod
    def log_and_return_error(
        logger_instance, error_msg: str, exception: Exception = None, return_value=None
    ):
        """
        记录错误并返回指定值

        Args:
            logger_instance: 日志记录器实例
            error_msg: 错误消息
            exception: 异常对象（可选）
            return_value: 要返回的值

        Returns:
            return_value参数指定的值
        """
        if exception:
            logger_instance.error(f"{error_msg}: {str(exception)}", exc_info=True)
        else:
            logger_instance.error(error_msg)
        return return_value

    @staticmethod
    def handle_connection_error(
        connection_id: str, operation: str, exception: Exception
    ):
        """
        处理连接相关错误

        Args:
            connection_id: 连接ID
            operation: 操作名称
            exception: 异常对象
        """
        error_msg = f"连接 {connection_id} 执行 {operation} 操作时出错"
        logger.error(f"{error_msg}: {str(exception)}", exc_info=True)

        # 可以在这里添加更多错误处理逻辑，如：
        # - 发送错误通知
        # - 清理无效连接
        # - 记录错误统计

    @staticmethod
    def handle_api_error(endpoint: str, request_data: dict, exception: Exception):
        """
        处理API相关错误

        Args:
            endpoint: API端点
            request_data: 请求数据
            exception: 异常对象
        """
        error_msg = f"API端点 {endpoint} 处理请求时出错"
        logger.error(f"{error_msg}: {str(exception)}")
        logger.debug(f"请求数据: {request_data}")


class SystemMonitor:
    """系统监控类，用于记录系统状态和性能指标"""

    def __init__(self):
        self.start_time = datetime.now()
        self.message_count = 0
        self.error_count = 0
        self.connection_events = []

    def record_message(self):
        """记录消息处理"""
        self.message_count += 1

    def record_error(self):
        """记录错误"""
        self.error_count += 1

    def record_connection_event(self, event_type: str, connection_id: str):
        """
        记录连接事件

        Args:
            event_type: 事件类型 ('connect', 'disconnect', 'error')
            connection_id: 连接ID
        """
        event = {
            "timestamp": datetime.now(),
            "type": event_type,
            "connection_id": connection_id,
        }
        self.connection_events.append(event)

        # 保持最近1000个事件
        if len(self.connection_events) > 1000:
            self.connection_events = self.connection_events[-1000:]

    def get_system_stats(self) -> Dict[str, Any]:
        """获取系统统计信息"""
        uptime = datetime.now() - self.start_time
        return {
            "uptime_seconds": int(uptime.total_seconds()),
            "uptime_formatted": str(uptime),
            "total_messages": self.message_count,
            "total_errors": self.error_count,
            "active_connections": connection_manager.get_connection_count(),
            "recent_events": len(self.connection_events),
        }

    def log_periodic_stats(self):
        """定期记录系统统计信息"""
        stats = self.get_system_stats()
        logger.info("=== 系统状态统计 ===")
        logger.info(f"运行时间: {stats['uptime_formatted']}")
        logger.info(f"处理消息总数: {stats['total_messages']}")
        logger.info(f"错误总数: {stats['total_errors']}")
        logger.info(f"当前活跃连接: {stats['active_connections']}")
        logger.info("==================")


# 创建全局错误处理器和系统监控器
error_handler = ErrorHandler()
system_monitor = SystemMonitor()

# 创建日志记录器
logger = logging.getLogger(__name__)


# ==================== 连接管理器 ====================


class ConnectionManager:
    """WebSocket连接管理器，负责管理所有活跃的WebSocket连接"""

    def __init__(self):
        """初始化连接管理器"""
        self.connections: Dict[str, WebSocketConnection] = {}  # 连接注册表
        self.lock = threading.Lock()  # 线程锁，保证线程安全
        logger.info("连接管理器初始化完成")

    def register_connection(self, flow: Any) -> str:
        """
        注册新的WebSocket连接

        Args:
            flow: mitmproxy的流对象

        Returns:
            str: 生成的连接ID
        """
        connection_id = str(uuid.uuid4())  # 生成唯一连接ID

        # 提取客户端和服务端地址信息
        client_address = f"{flow.client_conn.address[0]}:{flow.client_conn.address[1]}"
        server_address = (
            f"{flow.server_conn.address[0]}:{flow.server_conn.address[1]}"
            if flow.server_conn
            else "unknown"
        )

        # 创建连接对象
        connection = WebSocketConnection(
            connection_id=connection_id,
            client_address=client_address,
            server_address=server_address,
            created_at=datetime.now(),
            flow=flow,
        )

        # 线程安全地添加连接
        with self.lock:
            # 检查连接数限制
            if len(self.connections) >= config.max_connections:
                logger.warning(f"连接数已达到最大限制 {config.max_connections}")
                raise Exception(
                    f"Maximum connections limit ({config.max_connections}) reached"
                )

            self.connections[connection_id] = connection
            logger.info(
                f"新连接已注册: {connection_id}, 客户端: {client_address}, 服务端: {server_address}"
            )

        return connection_id

    def unregister_connection(self, connection_id: str) -> bool:
        """
        注销WebSocket连接

        Args:
            connection_id: 连接ID

        Returns:
            bool: 是否成功注销
        """
        with self.lock:
            if connection_id in self.connections:
                connection = self.connections[connection_id]
                connection.is_active = False
                del self.connections[connection_id]
                logger.info(f"连接已注销: {connection_id}")
                return True
            else:
                logger.warning(f"尝试注销不存在的连接: {connection_id}")
                return False

    def get_connection(self, connection_id: str) -> Optional[WebSocketConnection]:
        """
        获取指定的WebSocket连接

        Args:
            connection_id: 连接ID

        Returns:
            Optional[WebSocketConnection]: 连接对象，如果不存在则返回None
        """
        with self.lock:
            return self.connections.get(connection_id)

    def get_active_connections(self) -> List[Dict[str, Any]]:
        """
        获取所有活跃连接的信息

        Returns:
            List[Dict[str, Any]]: 连接信息列表
        """
        with self.lock:
            return [
                conn.to_dict() for conn in self.connections.values() if conn.is_active
            ]

    def send_message(
        self, connection_id: str, message: str, message_type: str = "text"
    ) -> bool:
        """
        向指定连接发送消息

        Args:
            connection_id: 目标连接ID
            message: 要发送的消息内容
            message_type: 消息类型 ('text' 或 'binary')

        Returns:
            bool: 是否成功发送
        """
        connection = self.get_connection(connection_id)
        if not connection or not connection.is_active:
            logger.error(f"连接不存在或已断开: {connection_id}")
            return False

        try:
            # 验证消息内容
            if not message:
                logger.error(f"消息内容为空，无法发送到连接 {connection_id}")
                return False

            # 处理消息内容转换
            message_content = self._prepare_message_content(message, message_type)
            if message_content is None:
                return False

            # 通过mitmproxy的流对象发送消息到客户端
            if hasattr(connection.flow, "websocket") and connection.flow.websocket:
                # 创建WebSocket消息对象
                ws_message = self._create_websocket_message(
                    message_content, message_type
                )

                # 将消息注入到WebSocket流中
                # connection.flow.websocket.messages.append(ws_message)
                master.mitmproxy_ctx.master.commands.call("inject.websocket", connection.flow, True, message_content, True)
                # 记录消息发送日志
                display_message = (
                    message[:100] + "..." if len(message) > 100 else message
                )
                logger.info(
                    f"消息已路由到连接 {connection_id} [{message_type}]: {display_message}"
                )

                return True
            else:
                logger.error(f"连接 {connection_id} 的WebSocket流无效或未建立")
                return False

        except Exception as e:
            logger.error(f"发送消息到连接 {connection_id} 时出错: {str(e)}")
            return False

    def _prepare_message_content(self, message: str, message_type: str):
        """
        准备消息内容，进行必要的格式转换

        Args:
            message: 原始消息内容
            message_type: 消息类型

        Returns:
            处理后的消息内容，失败时返回None
        """
        try:
            if message_type == "text":
                # 文本消息：确保是UTF-8编码的字节串
                if isinstance(message, str):
                    return message.encode("utf-8")
                elif isinstance(message, bytes):
                    # 验证是否为有效的UTF-8
                    message.decode("utf-8")
                    return message
                else:
                    logger.error(f"文本消息类型错误: {type(message)}")
                    return None

            elif message_type == "binary":
                # 二进制消息：转换为字节串
                if isinstance(message, str):
                    return message.encode("utf-8")
                elif isinstance(message, bytes):
                    return message
                else:
                    logger.error(f"二进制消息类型错误: {type(message)}")
                    return None
            else:
                logger.error(f"不支持的消息类型: {message_type}")
                return None

        except UnicodeDecodeError as e:
            logger.error(f"消息编码错误: {str(e)}")
            return None
        except Exception as e:
            logger.error(f"准备消息内容时出错: {str(e)}")
            return None

    def _create_websocket_message(self, content: bytes, message_type: str):
        """
        创建WebSocket消息对象

        Args:
            content: 消息内容（字节串）
            message_type: 消息类型

        Returns:
            WebSocket消息对象
        """
        if message_type == "text":
            return websocket.WebSocketMessage(
                type=1,  # TEXT
                content=content,
                from_client=False,  # 表示消息来自服务端发往客户端
                injected=True,      # 表示这是注入的消息
            )
        else:  # binary
            return websocket.WebSocketMessage(
                type=2,  # BINARY
                content=content,
                from_client=False,
                injected=True,      # 表示这是注入的消息
            )

    def broadcast_message(
        self,
        message: str,
        message_type: str = "text",
        exclude_connections: List[str] = None,
    ) -> Dict[str, bool]:
        """
        向所有活跃连接广播消息

        Args:
            message: 要广播的消息内容
            message_type: 消息类型
            exclude_connections: 要排除的连接ID列表

        Returns:
            Dict[str, bool]: 每个连接的发送结果
        """
        exclude_connections = exclude_connections or []
        results = {}

        with self.lock:
            active_connections = [
                conn_id
                for conn_id, conn in self.connections.items()
                if conn.is_active and conn_id not in exclude_connections
            ]

        for connection_id in active_connections:
            results[connection_id] = self.send_message(
                connection_id, message, message_type
            )

        success_count = sum(1 for success in results.values() if success)
        logger.info(f"广播消息完成: {success_count}/{len(results)} 连接成功接收")

        return results

    def get_connection_count(self) -> int:
        """
        获取当前活跃连接数

        Returns:
            int: 活跃连接数
        """
        with self.lock:
            return len([conn for conn in self.connections.values() if conn.is_active])


# 创建全局连接管理器实例
connection_manager = ConnectionManager()


# ==================== WebSocket拦截器 ====================


class WebSocketInterceptor:
    """
    mitmproxy WebSocket拦截器
    负责拦截WebSocket连接的建立、消息传输和关闭事件
    """

    def __init__(self):
        """初始化WebSocket拦截器"""
        self.flow_to_connection: Dict[Any, str] = {}  # 流对象到连接ID的映射
        logger.info("WebSocket拦截器初始化完成")

    def websocket_start(self, flow):
        """
        WebSocket连接建立时的处理函数

        Args:
            flow: mitmproxy的流对象
        """
        try:
            # 在连接管理器中注册新连接
            connection_id = connection_manager.register_connection(flow)

            # 建立流对象到连接ID的映射
            self.flow_to_connection[flow] = connection_id

            # 记录连接建立事件
            client_addr = f"{flow.client_conn.address[0]}:{flow.client_conn.address[1]}"
            server_addr = (
                f"{flow.server_conn.address[0]}:{flow.server_conn.address[1]}"
                if flow.server_conn
                else "unknown"
            )

            logger.info(f"WebSocket连接已建立 - ID: {connection_id}")
            logger.info(f"  客户端: {client_addr}")
            logger.info(f"  服务端: {server_addr}")
            logger.info(f"  URL: {flow.request.pretty_url}")
            logger.info(
                f"  当前活跃连接数: {connection_manager.get_connection_count()}"
            )

            # 记录连接事件
            system_monitor.record_connection_event("connect", connection_id)

        except Exception as e:
            error_handler.handle_connection_error("unknown", "建立连接", e)
            system_monitor.record_error()

    def websocket_message(self, flow):
        """
        WebSocket消息传输时的处理函数

        Args:
            flow: mitmproxy的流对象
        """
        try:
            # 获取连接ID
            connection_id = self.flow_to_connection.get(flow)
            if not connection_id:
                logger.warning("收到消息但找不到对应的连接ID")
                return

            # 获取最新的消息
            # while True:
            if flow.websocket and flow.websocket.messages:
                latest_message = flow.websocket.messages[-1]

                # 确定消息方向
                direction = (
                    "客户端 -> 服务端"
                    if latest_message.from_client
                    else "服务端 -> 客户端"
                )

                # 获取消息内容（截断长消息用于日志）
                content = latest_message.content
                if isinstance(content, bytes):
                    try:
                        content_str = content.decode("utf-8")
                    except UnicodeDecodeError:
                        content_str = f"<binary data: {len(content)} bytes>"
                else:
                    content_str = str(content)

                # 截断过长的消息内容用于日志显示
                display_content = (
                    content_str[:200] + "..." if len(content_str) > 200 else content_str
                )

                # 记录消息传输
                logger.debug(
                    f"WebSocket消息 [{connection_id}] {direction}: {display_content}"
                )

                # 记录消息处理统计
                system_monitor.record_message()

                # 这里消息会自动转发，mitmproxy会处理透明代理
                # 我们只需要记录和监控，不需要手动转发
                # else:
                #     break
        except Exception as e:
            logger.error(f"处理WebSocket消息时出错: {str(e)}")

    def websocket_end(self, flow):
        """
        WebSocket连接关闭时的处理函数

        Args:
            flow: mitmproxy的流对象
        """
        try:
            # 获取连接ID
            connection_id = self.flow_to_connection.get(flow)
            if connection_id:
                # 从连接管理器中注销连接
                connection_manager.unregister_connection(connection_id)

                # 清理流对象映射
                del self.flow_to_connection[flow]

                logger.info(f"WebSocket连接已关闭 - ID: {connection_id}")
                logger.info(
                    f"  当前活跃连接数: {connection_manager.get_connection_count()}"
                )

                # 记录断开连接事件
                system_monitor.record_connection_event("disconnect", connection_id)
            else:
                logger.warning("WebSocket连接关闭但找不到对应的连接ID")

        except Exception as e:
            connection_id = self.flow_to_connection.get(flow, "unknown")
            error_handler.handle_connection_error(connection_id, "关闭连接", e)
            system_monitor.record_error()

    def get_connection_id_by_flow(self, flow) -> Optional[str]:
        """
        根据流对象获取连接ID

        Args:
            flow: mitmproxy的流对象

        Returns:
            Optional[str]: 连接ID，如果不存在则返回None
        """
        return self.flow_to_connection.get(flow)


# 创建全局WebSocket拦截器实例
websocket_interceptor = WebSocketInterceptor()


# ==================== HTTP API服务器 ====================


class HTTPAPIServer:
    """
    HTTP API服务器
    提供RESTful接口用于外部系统与WebSocket连接交互
    """

    def __init__(self):
        """初始化HTTP API服务器"""
        self.app = web.Application()
        self.setup_routes()
        logger.info("HTTP API服务器初始化完成")

    def setup_routes(self):
        """设置API路由"""
        self.app.router.add_post("/send", self.send_message)
        self.app.router.add_get("/connections", self.get_connections)
        self.app.router.add_delete(
            "/connections/{connection_id}", self.disconnect_connection
        )
        self.app.router.add_get("/health", self.health_check)

        # 添加CORS支持
        self.app.middlewares.append(self.cors_middleware)

    @web.middleware
    async def cors_middleware(self, request, handler):
        """CORS中间件，允许跨域请求"""
        response = await handler(request)
        response.headers["Access-Control-Allow-Origin"] = "*"
        response.headers["Access-Control-Allow-Methods"] = "GET, POST, DELETE, OPTIONS"
        response.headers["Access-Control-Allow-Headers"] = "Content-Type"
        return response

    async def send_message(self, request):
        """
        POST /send - 发送消息到指定WebSocket连接

        请求体格式:
        {
            "connection_id": "连接ID",
            "message": "消息内容",
            "message_type": "text"  // 可选，默认为text
        }
        """
        try:
            # 解析请求体
            try:
                data = await request.json()
            except json.JSONDecodeError:
                return web.json_response({"error": "Invalid JSON format"}, status=400)

            # 创建并验证请求对象
            send_request = SendMessageRequest.from_dict(data)
            errors = send_request.validate()

            if errors:
                return web.json_response(
                    {"error": "Validation failed", "details": errors}, status=400
                )

            # 发送消息
            success = connection_manager.send_message(
                send_request.connection_id,
                send_request.message,
                send_request.message_type,
            )

            if success:
                logger.info(f"HTTP API: 消息已发送到连接 {send_request.connection_id}")
                return web.json_response(
                    {
                        "success": True,
                        "message": "Message sent successfully",
                        "connection_id": send_request.connection_id,
                    }
                )
            else:
                return web.json_response(
                    {
                        "error": "Failed to send message",
                        "connection_id": send_request.connection_id,
                    },
                    status=404,
                )

        except Exception as e:
            logger.error(f"HTTP API发送消息时出错: {str(e)}")
            return web.json_response(
                {"error": "Internal server error", "details": str(e)}, status=500
            )

    async def get_connections(self, request):
        """
        GET /connections - 获取所有活跃WebSocket连接信息
        """
        try:
            connections = connection_manager.get_active_connections()

            return web.json_response(
                {"success": True, "count": len(connections), "connections": connections}
            )

        except Exception as e:
            logger.error(f"HTTP API获取连接列表时出错: {str(e)}")
            return web.json_response(
                {"error": "Internal server error", "details": str(e)}, status=500
            )

    async def disconnect_connection(self, request):
        """
        DELETE /connections/{connection_id} - 断开指定WebSocket连接
        """
        try:
            connection_id = request.match_info["connection_id"]

            # 获取连接信息
            connection = connection_manager.get_connection(connection_id)
            if not connection:
                return web.json_response(
                    {"error": "Connection not found", "connection_id": connection_id},
                    status=404,
                )

            # 断开连接
            success = connection_manager.unregister_connection(connection_id)

            if success:
                logger.info(f"HTTP API: 连接 {connection_id} 已断开")
                return web.json_response(
                    {
                        "success": True,
                        "message": "Connection disconnected successfully",
                        "connection_id": connection_id,
                    }
                )
            else:
                return web.json_response(
                    {
                        "error": "Failed to disconnect connection",
                        "connection_id": connection_id,
                    },
                    status=500,
                )

        except Exception as e:
            logger.error(f"HTTP API断开连接时出错: {str(e)}")
            return web.json_response(
                {"error": "Internal server error", "details": str(e)}, status=500
            )

    async def health_check(self, request):
        """
        GET /health - 健康检查端点
        """
        return web.json_response(
            {
                "status": "healthy",
                "active_connections": connection_manager.get_connection_count(),
                "timestamp": datetime.now().isoformat(),
            }
        )

    async def start_server(self):
        """启动HTTP API服务器"""
        runner = web.AppRunner(self.app)
        await runner.setup()

        site = web.TCPSite(runner, config.host, config.api_port)
        await site.start()

        logger.info(f"HTTP API服务器已启动: http://{config.host}:{config.api_port}")
        logger.info("可用的API端点:")
        logger.info(f"  POST   http://{config.host}:{config.api_port}/send")
        logger.info(f"  GET    http://{config.host}:{config.api_port}/connections")
        logger.info(
            f"  DELETE http://{config.host}:{config.api_port}/connections/{{id}}"
        )
        logger.info(f"  GET    http://{config.host}:{config.api_port}/health")


# 创建全局HTTP API服务器实例
http_api_server = HTTPAPIServer()


# ==================== 主程序入口和配置 ====================


async def start_http_server():
    """启动HTTP API服务器"""
    try:
        await http_api_server.start_server()
    except Exception as e:
        logger.error(f"启动HTTP API服务器失败: {str(e)}")
        raise


async def start_mitmproxy():
    """启动mitmproxy服务器"""
    try:
        # 配置mitmproxy选项
        opts = Options(
            listen_port=config.proxy_port,
            listen_host=config.host,
        )

        # 创建DumpMaster实例
        master = DumpMaster(opts)

        # 添加WebSocket拦截器插件
        master.addons.add(websocket_interceptor)

        logger.info(f"mitmproxy服务器启动: http://{config.host}:{config.proxy_port}")
        logger.info("请配置客户端使用以下代理设置:")
        logger.info(f"  HTTP代理: {config.host}:{config.proxy_port}")
        logger.info(f"  HTTPS代理: {config.host}:{config.proxy_port}")

        # 运行mitmproxy
        await master.run()

    except Exception as e:
        logger.error(f"启动mitmproxy服务器失败: {str(e)}")
        raise


async def periodic_stats_logger():
    """定期记录系统统计信息"""
    while True:
        try:
            await asyncio.sleep(300)  # 每5分钟记录一次
            system_monitor.log_periodic_stats()
        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.error(f"记录统计信息时出错: {str(e)}")


async def main_async():
    """异步主函数"""
    try:
        logger.info("正在启动WebSocket代理桥接服务...")

        # 启动HTTP API服务器
        await start_http_server()

        # 启动定期统计日志任务
        stats_task = asyncio.create_task(periodic_stats_logger())

        logger.info("所有服务已启动，系统运行中...")
        logger.info("按 Ctrl+C 停止服务")

        # 启动mitmproxy服务器任务
        proxy_task = asyncio.create_task(start_mitmproxy())

        # 等待中断信号
        try:
            while True:
                await asyncio.sleep(1)
        except KeyboardInterrupt:
            logger.info("收到中断信号，正在关闭服务...")
            stats_task.cancel()
            proxy_task.cancel()

    except Exception as e:
        logger.error(f"主程序运行时出错: {str(e)}")
        raise


def load_config_file(config_path: str = "config.json"):
    """
    从配置文件加载配置

    Args:
        config_path: 配置文件路径
    """
    try:
        import os

        if os.path.exists(config_path):
            with open(config_path, "r", encoding="utf-8") as f:
                config_data = json.load(f)

            # 更新配置
            for key, value in config_data.items():
                if hasattr(config, key):
                    setattr(config, key, value)
                    logger.info(f"从配置文件加载: {key} = {value}")
                else:
                    logger.warning(f"未知的配置项: {key}")
        else:
            logger.info(f"配置文件 {config_path} 不存在，使用默认配置")
    except Exception as e:
        logger.error(f"加载配置文件失败: {str(e)}")


def create_default_config_file(config_path: str = "config.json"):
    """
    创建默认配置文件

    Args:
        config_path: 配置文件路径
    """
    try:
        default_config = {
            "proxy_port": 8999,
            "api_port": 8998,
            "log_level": "INFO",
            "max_connections": 100,
            "host": "127.0.0.1",
        }

        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(default_config, f, indent=2, ensure_ascii=False)

        logger.info(f"已创建默认配置文件: {config_path}")
    except Exception as e:
        logger.error(f"创建配置文件失败: {str(e)}")


def signal_handler(signum, frame):
    """信号处理器，用于优雅关闭"""
    logger.info(f"收到信号 {signum}，正在关闭服务...")
    import sys

    sys.exit(0)


if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(
        description="WebSocket代理桥接服务 - 使用mitmproxy拦截WebSocket并提供HTTP控制接口",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  python websocket_proxy_bridge.py                    # 使用默认配置启动
  python websocket_proxy_bridge.py --proxy-port 8080  # 指定代理端口
  python websocket_proxy_bridge.py --config config.json # 使用配置文件
  python websocket_proxy_bridge.py --create-config    # 创建默认配置文件

API使用示例:
  # 获取连接列表
  curl http://127.0.0.1:8998/connections
  
  # 发送消息
  curl -X POST http://127.0.0.1:8998/send \\
    -H "Content-Type: application/json" \\
    -d '{"connection_id": "连接ID", "message": "Hello WebSocket!"}'
        """,
    )

    parser.add_argument(
        "--proxy-port", type=int, default=8999, help="mitmproxy监听端口 (默认: 8999)"
    )
    parser.add_argument(
        "--api-port", type=int, default=8998, help="HTTP API端口 (默认: 8998)"
    )
    parser.add_argument(
        "--host", default="127.0.0.1", help="服务器绑定地址 (默认: 127.0.0.1)"
    )
    parser.add_argument(
        "--log-level",
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="日志级别 (默认: INFO)",
    )
    parser.add_argument(
        "--max-connections", type=int, default=100, help="最大连接数 (默认: 100)"
    )
    parser.add_argument("--config", type=str, help="配置文件路径")
    parser.add_argument(
        "--create-config", action="store_true", help="创建默认配置文件并退出"
    )

    args = parser.parse_args()

    # 如果只是创建配置文件
    if args.create_config:
        create_default_config_file(args.config or "config.json")
        exit(0)

    # 加载配置文件（如果指定）
    if args.config:
        load_config_file(args.config)

    # 命令行参数覆盖配置文件
    config.proxy_port = args.proxy_port
    config.api_port = args.api_port
    config.host = args.host
    config.log_level = args.log_level
    config.max_connections = args.max_connections

    # 设置日志
    setup_logging()

    # 注册信号处理器
    import signal

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # 显示启动信息
    logger.info("=" * 60)
    logger.info("WebSocket代理桥接服务启动")
    logger.info("=" * 60)
    logger.info(f"代理端口: {config.proxy_port}")
    logger.info(f"API端口: {config.api_port}")
    logger.info(f"绑定地址: {config.host}")
    logger.info(f"日志级别: {config.log_level}")
    logger.info(f"最大连接数: {config.max_connections}")
    logger.info("=" * 60)

    try:
        # 运行主程序
        asyncio.run(main_async())
    except KeyboardInterrupt:
        logger.info("服务已停止")
    except Exception as e:
        logger.error(f"程序异常退出: {str(e)}")
        exit(1)
